from django.contrib.auth import get_user_model
from django.db import transaction
from rest_framework.exceptions import AuthenticationFailed
from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.tokens import AccessToken

from user.models import Profile, GroupExtraInfo


class CustomJWTAuthentication(JWTAuthentication):
    def authenticate(self, request):
        token = request.headers.get("Authorization")
        if token:
            try:
                # 去掉 "Bearer " 前缀
                access_token = token.split(" ")[1]
                # 解析 token
                validated_token = AccessToken(access_token)
                username = validated_token.get("username")
                nickname = validated_token.get("nickname")
                if not username:
                    raise AuthenticationFailed("Username is required in the token.")
                login_way = validated_token.get('type')

                # 如果从本系统过来
                if login_way != "sso":
                    # 尝试查找用户
                    try:
                        User = get_user_model()
                        user = User.objects.get(username=username)
                    except User.DoesNotExist:
                        raise AuthenticationFailed(f"User '{username}' not found.")

                    return (user, validated_token)

                # 从统一认证系统（SSO）过来
                with transaction.atomic():
                    try:
                        User = get_user_model()
                        user = User.objects.get(username=username)
                    except User.DoesNotExist:
                        # 用户不存在，创建新用户
                        user = User.objects.create_user(
                            username=username, password="12345678"
                        )
                        Profile.objects.create(
                            user=user, chinese_name=nickname
                        )

                    # 处理用户的角色
                    group_codenames = validated_token.get('roles', [])
                    for group_codename in group_codenames:
                        try:
                            group_extra_info = GroupExtraInfo.objects.get(
                                code_name=group_codename
                            )
                            group = group_extra_info.group
                            group.user_set.add(user)
                        except GroupExtraInfo.DoesNotExist:
                            # raise AuthenticationFailed("Group not found.")
                            return {'message': "Group not found.", 'code': 0, 'data': ''}

                    return (user, validated_token)

            except Exception as e:
                raise AuthenticationFailed(f"Error parsing token: {str(e)}")
        else:
            return None
